The PR description says that this is an experimental feature and is disabled by default. Is it intentional that it is enabled by default here?

These changes seem specific to the ignored test, but will impact all of the existing tests?

I don't see any new tests for the functional changes in this PR. Could you add functional tests that use small amounts of data, just to test for correctness?

The entire CI is passed with useLargeDataTypes enabled.
on CI grade machines we cannot reproduce issue as we run out of memory earlier than hit the offset limit

[P1] Keep the final aggregate’s input and spill schemas consistent with promoted keys. This promotion also runs in Final mode, while its child still produces Utf8/Binary. DataFusion 55.1.0 uses agg.input().schema() as the final aggregate’s state schema. When memory pressure triggers spilling, it stamps that small-offset schema onto promoted group arrays and fails with expected Utf8 but found LargeUtf8 (likewise for binary). The output alignment cannot fix an error inside the aggregate. With the default enabled, ordinary spilling GROUP BY/DISTINCT queries now fail. Promote the final input through a matching projection, or otherwise make its spill state schema agree with the promoted keys.

Evidence: Bounded reproduction in /tmp/comet-4791-validation/src/main.rs: native Final distinct aggregation over 80 batches of 256 rows, 10,000 distinct 128-byte keys, batch size 256, and a 1 MiB FairSpillPool. Without promotion, both Utf8 and Binary return 10,000 keys after five spills. With the PR’s grouping-expression promotion, both fail before returning rows with the corresponding small/large type mismatch. Promoting the input through ProjectionExec restores 10,000 results and five spills. DataFusion’s aggregate_hash_table/final_table.rs takes the state schema from agg.input().schema(), and common.rs::take_state_batch constructs the failing batch.

[P2] Preserve aggregate metrics when installing this wrapper. SchemaAlignExec has no metrics() implementation, but it replaces AggregateExec as the native plan passed to SparkPlan::new. to_native_metric_node consequently reads no metrics and never visits the wrapped aggregate. Every promoted string/binary aggregation loses output-row, execution-time, and spill reporting, including the spill information propagated into Spark task metrics. Forward the aggregate’s metrics through the wrapper or explicitly retain it as the metric source.

Evidence: The exact-source harness executes string and binary aggregates containing duplicates, nulls, and empty keys. Both return four groups, and the underlying AggregateExec reports output_rows=4 and nonzero elapsed_compute, while the wrapper returns metrics=None. SparkPlan::new leaves additional_native_plans empty, and native/core/src/execution/metrics/utils.rs::to_native_metric_node only reads the root metrics in that case.

I can confirm this end to end with a full build. SELECT k, sum(v) FROM t GROUP BY k on a string key shows the partial CometHashAggregateExec with output_rows=0 and elapsed_compute=0 when the flag is on, against 20000 rows and about 50 ms with it off. The final aggregate's elapsed_compute drops from about 10 ms to 17 µs. When fixing it, note that to_native_metric_node skips output_rows from additional_native_plans, so registering the AggregateExec through SparkPlan::new_with_additional also needs the wrapper to record its own output_rows. A test asserting these metrics on a string-keyed aggregate would catch this. The one in CometExecSuite uses a numeric key.

[P2] Avoid copying value buffers when the offsets already fit. This builder path runs for every promoted string key, with the equivalent copy for binary, even when no splitting or rebasing is needed. Since promotion defaults to enabled, each partial/final aggregation adds a full allocation and copy of its emitted key bytes. Arrow can narrow ordinary batches while sharing those bytes. Use that fast path when offsets fit, reserving rebuilding for overflow slices, or rebase offsets against a shared value-buffer slice.

Evidence: The exact SchemaAlignExec source copies a 4 MiB value buffer for 8,192 keys of 512 bytes, verified by buffer-pointer comparison. Arrow 59.3.0’s cast shares it and produces identical values. An optimized benchmark of the PR’s exact builder branch measured median conversion times of 168.1 µs versus 8.4 µs for that batch, and 7.81 ms versus 8.3 µs for 8,192 keys of 8 KiB. These are narrowing-only measurements, not whole-query timings. Reproduction: /tmp/comet-4791-cast-bench/src/main.rs.

It looks like the rebase brought back the ColumnAction enum that #5138 removed from this file. #5138 routed SchemaAlignStream through cast_and_stamp_schema, so every batch is checked against the target schema and a failed cast names the operator and the column path. Here the choice is made once at plan time. Passthrough columns are no longer checked per batch, and the Cast arm returns a bare arrow error. This operator sits in front of every native shuffle writer, so that applies with the flag off too. Could the Large-to-small columns be handled first, with the result then handed to cast_and_stamp_schema?

With the flag on, reusing SchemaAlignExec above the aggregate hits the warning just above and logs ShuffleWriter input schema mismatch on col[0] 'col_0': child produced LargeUtf8, catalyst declared Utf8. Inserting a cast; please file the upstream function bug at .../issues/4515 on each executor. I saw it in my local runs. There is no shuffle writer or upstream bug involved here, and the module doc says this operator is enclosed by shuffle on purpose. Would a small operator dedicated to the cast-back be cleaner? It could record its own metrics, which would also help with the metrics problem sunchao raised in planner.rs.

I asked in July whether defaulting this to true was intentional, and the doc string still says it defaults to false. I measured it on a release build. SELECT count(*), sum(s) FROM (SELECT k, sum(v) s FROM t GROUP BY k) over 8M rows with 4M distinct 21-byte keys took 494 ms at the median with the flag on and 450 ms with it off, about 10% slower over five alternating runs. Low-cardinality and 200-byte keys were within noise. My laptop was busy, so treat the numbers as rough. The overflow needs more than 2 GiB of distinct key bytes in one task, and apache/datafusion#24704 tracks the real fix. Could this default to false, with the tuning guide pointing at it for the offset overflow error? Either way the doc needs updating. It describes the cast-back as a Projection and says the overhead is O(rows), but as sunchao pointed out, the cast-back copies every key byte.

The other aggregate keys use spark.comet.exec.aggregate.*, and config_conventions.md asks for feature flags to end in .enabled. This key will be covered by the versioning policy once it ships, so could we settle on something like spark.comet.exec.aggregate.largeGroupKeys.enabled now?

Since the aggregate now casts back before anything above it sees the batch, I don't think this feature can put LargeUtf8 into a shuffle block. If one did, the remote read path would reject it, because remote_schema.rs treats LargeUtf8 against Utf8 as an incompatible type. On the local path, ShuffleScanStream::poll_next already reconciles every column through cast_and_stamp_schema, so the extra cast here duplicates that. Could this go back to unpacking dictionaries only? Returning an error instead of the old expect is a good change and worth keeping.

Which path needs getFieldVector to accept LargeVarCharVector now that the aggregate casts back natively? As far as I can tell, the only Large vectors on the JVM side come from PyArrow UDFs returning large_string. So this changes native C2R export and broadcast serialization for that path without a test. It also makes the isSupportedFieldVector doc, the isArrowBacked comment and the UtilsSuite test comment wrong, since all three assume getFieldVector rejects these vectors. Could this be dropped here, or moved to its own PR with a PyArrow test? The same question applies to the LargeUtf8 passthrough at columnar_to_row.rs:1041.

I can't find a way to reach these new arms, or the matching ones in cast.rs. types.proto has no Large type ids, so the serde never emits a Cast to or from one. The Parquet schema adapter's casts already go through the is_adapting_schema branch in cast_array. The comment also says SchemaAlignExec pre-splits arrays before they get here, but SchemaAlignExec builds its arrays directly and never calls this cast. Could these arms be removed from this PR?

This test is ignored, so CI never runs it. The sparkConf comment at line 60 still describes the off-heap and memory pool settings that 7252adb removed. This test now sets the pool through withSQLConf, which that comment says won't take effect. You mentioned in July that CI passes with the flag on. That covers the common path, but nothing exercises the split in compute_row_ranges or checks the metrics. Could the byte cap be a parameter so a Rust unit test can split a small batch? And could this test and the stale comment be replaced with small tests that run? If the default moves to false, those tests would also need to turn the flag on explicitly.